import numpy as np
import pandas as pd
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.model_selection import GridSearchCV
from sklearn.preprocessing import LabelEncoder,OneHotEncoder,StandardScaler,OrdinalEncoder,LabelBinarizer
from sklearn.metrics import classification_report, accuracy_score, precision_score, recall_score, f1_score, \
    confusion_matrix
from sklearn.model_selection import train_test_split,cross_val_score
from sklearn.neighbors import KNeighborsClassifier
import matplotlib.pyplot as plt

import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

# 定义三个模型的名称和三个不同指标的得分
models = ['KNN', 'Bayesian', 'RNN']
precision = [0.87, 0.80, 0.85]
recall = [0.87, 0.77, 0.85]
f1_score = [0.87, 0.78, 0.85]

# 将得分组合成一个 2D 数组
scores = np.array([precision, recall, f1_score])

# 绘制堆叠条形图
x = np.arange(len(models))
width = 0.35
fig, ax = plt.subplots()
ax.bar(x, scores[0], width, label='Precision')
ax.bar(x + width, scores[1], width, label='Recall')
ax.bar(x + width*2, scores[2], width, label='F1-score')

# 设置轴标签、标题和图例
ax.set_xlabel('Model')
ax.set_ylabel('Score')
ax.set_title('Model Comparison')
ax.set_xticks(x + width)
ax.set_xticklabels(models)
ax.legend()

# 显示图形
#plt.show()

# 热力图
import pandas as pd
import seaborn as sns
from sklearn.preprocessing import LabelEncoder

df1 = pd.read_csv('pdb_data_no_dups.csv')
df2 = pd.read_csv('pdb_data_seq.csv')

# 合并数据集
df_all = df1.set_index('structureId').merge(df2.set_index('structureId'), on='structureId', how='left')
# 从两个文件（df_seq和df_char）抽提出蛋白序列
protein_char = df1[df1.macromoleculeType == 'Protein']
protein_seq = df2[df2.macromoleculeType == 'Protein']

# 以‘structureId’作为合并的连接键，从两个表格里筛选需要的列进行数据合并
protein_char = protein_char[['structureId', 'classification']]
protein_seq = protein_seq[['structureId', 'sequence']]

# 以‘structureId’为连接键，先将‘structureId’设为index,再用join()合并
model_df = protein_seq.set_index('structureId').join(protein_char.set_index('structureId'))

# 删除含有缺失值的列
model_df = model_df.dropna()

# 统计classification分类信息
counts = model_df['classification'].value_counts()

# 选择数量超过1000的数据
types = np.asarray(counts[(counts > 1000)].index)

# 从model_df中抽取满足条件的数据（数量超1000）
data = model_df[model_df.classification.isin(types)]


X=data['sequence']
Y=data['classification']
X_train,X_test,Y_train,Y_test=train_test_split(X,Y,test_size = 0.2, random_state = 1)

vect=CountVectorizer(analyzer='char_wb',ngram_range=(4,4))
vect.fit(X_train)
X_train_df=vect.transform(X_train)
X_test_df=vect.transform(X_test)

from sklearn.naive_bayes import MultinomialNB
model=MultinomialNB()
model.fit(X_train_df,Y_train)
NB_pred=model.predict(X_test_df)

conf_mat=confusion_matrix(Y_test,NB_pred,labels=types)
conf_mat_Nor = conf_mat.astype('float') / conf_mat.sum(axis=1)[:, np.newaxis]
plt.figure(figsize=(13,8))
sns.heatmap(conf_mat)
plt.show()




